import torch

from utils.bbox_transform import cxcywh_to_xyxy


def iou_xyxy(gt_xyxy, pred_xyxy):
    """
    计算基于左上角和右下角坐标 (x_min, y_min, x_max, y_max) 的 IOU。
    输入:
        gt_xyxy: 张量，形状为 [N, 4]，包含 Ground Truth 的 `(x_min, y_min, x_max, y_max)`。
        pred_xyxy: 张量，形状为 [N, 4]，包含预测框的 `(x_min, y_min, x_max, y_max)`。
    输出:
        张量，形状为 [N]，每个值为对应框的 IOU。
    """
    inter_xmin = torch.max(gt_xyxy[:, 0], pred_xyxy[:, 0])
    inter_ymin = torch.max(gt_xyxy[:, 1], pred_xyxy[:, 1])
    inter_xmax = torch.min(gt_xyxy[:, 2], pred_xyxy[:, 2])
    inter_ymax = torch.min(gt_xyxy[:, 3], pred_xyxy[:, 3])
    inter_area = torch.clamp(inter_xmax - inter_xmin, min=0) * torch.clamp(inter_ymax - inter_ymin, min=0)

    gt_area = (gt_xyxy[:, 2] - gt_xyxy[:, 0]) * (gt_xyxy[:, 3] - gt_xyxy[:, 1])
    pred_area = (pred_xyxy[:, 2] - pred_xyxy[:, 0]) * (pred_xyxy[:, 3] - pred_xyxy[:, 1])
    union_area = gt_area + pred_area - inter_area

    return inter_area / union_area

def iou_cxcywh(gt_cxcywh, pred_cxcywh):
    """
    计算基于中心点和宽高 (cx, cy, w, h) 的 IOU。
    输入:
        gt_cxcywh: 张量，形状为 [N, 4]，包含 Ground Truth 的 `(cx, cy, w, h)`。
        pred_cxcywh: 张量，形状为 [N, 4]，包含预测框的 `(cx, cy, w, h)`。
    输出:
        张量，形状为 [N]，每个值为对应框的 IOU。
    """
    gt_xyxy = cxcywh_to_xyxy(gt_cxcywh)
    pred_xyxy = cxcywh_to_xyxy(pred_cxcywh)

    return iou_xyxy(gt_xyxy, pred_xyxy)
